package com.sjk.flink.kafka;

import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.util.Preconditions;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author cbry
 * @version 1.0.0
 * @ClassName FlinkRoundRobinPartitioner
 * @createTime 2022年06月02日 11:56:00
 * 轮询写入topic的分区，参考org.apache.kafka.producer.RoundRobinPartitioner
 * 以及（默认）org.apache.flink.streaming.connectors.kafka.partitioner.FlinkFixedPartitioner;
 *
 * @Desc 轮询的是并行度（默认机器CPU核数）子线程的写入。
 */
public class FlinkRoundRobinPartitioner<T> extends FlinkKafkaPartitioner<T> {

    private static final long serialVersionUID = -3785323215233858777L;
    private int parallelInstanceId;

    private final ConcurrentMap<String, AtomicInteger> topicCounterMap = new ConcurrentHashMap();

    public FlinkRoundRobinPartitioner() {
    }

    /**
     * 构建子线程创建时候执行，只执行一次相当于赋值线程名称标识：并行度
     */
    @Override
    public void open(int parallelInstanceId, int parallelInstances) {
        System.out.println("写入分区：open");
        Preconditions.checkArgument(parallelInstanceId >= 0, "Id of this subtask cannot be negative.");
        Preconditions.checkArgument(parallelInstances > 0, "Number of subtasks must be larger than 0.");
        this.parallelInstanceId = parallelInstanceId;
    }

    @Override
    public int partition(T record, byte[] key, byte[] value, String targetTopic, int[] partitions) {
        //并行度从1开始算 ， 分区从0开始算
        //int writePartiton = this.nextValue(targetTopic) % partitions.length;
        //System.out.println("并行度" + parallelInstanceId + "写入分区：" + writePartiton);
        return  this.nextValue(targetTopic) % partitions.length;
    }

    /**
     * 越界问题：Integer.MAX_VALUE + 1 =  Integer.MIX_VALUE;
     */
    private int nextValue(String topic) {
        AtomicInteger counter = (AtomicInteger)this.topicCounterMap.computeIfAbsent(topic, (k) -> {
            System.out.println("写入分区：初始化");
            return new AtomicInteger(0);
        });
        return counter.getAndIncrement();
    }
}